{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NOTE: This notebook contains interactive widgets that will work on Binder but not on Google Colab.  As such, the standard Colab setup script is not included."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from ipywidgets import interact\n",
    "from pydrake.all import MathematicalProgram, Solve, eq, le, ge\n",
    "\n",
    "from underactuated.jupyter import SetupMatplotlibBackend\n",
    "SetupMatplotlibBackend()\n",
    "\n",
    "prog = MathematicalProgram()\n",
    "\n",
    "lp_relaxation = False\n",
    "N = 20     # number of timesteps\n",
    "h = 0.1    # timestep\n",
    "bigM = 10 \n",
    "q_start = np.array([0., -2.])\n",
    "q_goal = np.array([0., 2.])\n",
    "\n",
    "q = prog.NewContinuousVariables(2, N, \"q\")   # positions\n",
    "v = prog.NewContinuousVariables(2, N, \"v\")   # velocities\n",
    "u = prog.NewContinuousVariables(2, N-1, \"u\") # acceleration command\n",
    "if lp_relaxation:\n",
    "  b = prog.NewContinuousVariables(4, N-1, \"b\") # binaries for collisions\n",
    "  prog.AddBoundingBoxConstraint(0, 1, b)\n",
    "else:\n",
    "  b = prog.NewBinaryVariables(4, N-1, \"b\")     # binaries for collisions\n",
    "\n",
    "# Initial value constraint:\n",
    "start = prog.AddLinearConstraint(eq(q[:,0], q_start))\n",
    "prog.AddLinearConstraint(eq(v[:,0], [0., 0.]))\n",
    "\n",
    "# Final value constraint:\n",
    "goal = prog.AddLinearConstraint(eq(q[:,N-1], q_goal))\n",
    "prog.AddLinearConstraint(eq(v[:,N-1], [0., 0.]))\n",
    "\n",
    "for n in range(0, N-1):\n",
    "  # Dynamics constraints:\n",
    "  prog.AddLinearConstraint(eq(q[:,n+1], q[:,n] + h*v[:,n]))\n",
    "  prog.AddLinearConstraint(eq(v[:,n+1], v[:,n] + h*u[:,n]))\n",
    "  \n",
    "  # Stage cost: sum u^2\n",
    "  prog.AddQuadraticCost(u[:,n].dot(u[:,n]))\n",
    "  \n",
    "  # Collision avoidance constraints:\n",
    "  prog.AddLinearConstraint(q[0,n+1] >= 1 - (1-b[0,n])*bigM)\n",
    "  prog.AddLinearConstraint(q[0,n+1] <= -1 + (1-b[1,n])*bigM)\n",
    "  prog.AddLinearConstraint(q[1,n+1] >= 1 - (1-b[2,n])*bigM)\n",
    "  prog.AddLinearConstraint(q[1,n+1] <= -1 + (1-b[3,n])*bigM)\n",
    "  prog.AddLinearConstraint(np.sum(b[:,n]) >= 1)\n",
    " \n",
    "# TODO(russt): This currently fails with only open-source solvers.  See drake #13299\n",
    "result = Solve(prog)\n",
    "  \n",
    "qt = result.GetSolution(q)\n",
    "bt = result.GetSolution(b)\n",
    "\n",
    "fig, ax = plt.subplots(2, 1)\n",
    "line, = ax[0].plot(qt[0,:], qt[1,:], 'b.-')\n",
    "start_pt, = ax[0].plot(q_start[0], q_start[1], 'b*', markersize=15)\n",
    "goal_pt, = ax[0].plot(q_goal[0], q_goal[1], 'g*', markersize=15)\n",
    "ax[0].fill([1, 1, -1, -1], [-1, 1, 1, -1], 'r')\n",
    "ax[0].axis('equal')\n",
    "ax[0].set_ylim(-4, 4)\n",
    "blines = ax[1].plot(bt.T)\n",
    "ax[1].set_ylim(-.1, 1.1)\n",
    "ax[1].set_xticks(range(0,N,2))\n",
    "ax[1].legend(blines, (\"right\", \"left\", \"top\", \"bottom\"))\n",
    "\n",
    "def resolve(start_x, start_y, goal_x, goal_y):\n",
    "  start.evaluator().UpdateLowerBound([start_x, start_y])\n",
    "  start.evaluator().UpdateUpperBound([start_x, start_y])\n",
    "  goal.evaluator().UpdateLowerBound([goal_x, goal_y])\n",
    "  goal.evaluator().UpdateUpperBound([goal_x, goal_y])\n",
    "\n",
    "  result = Solve(prog)\n",
    "  print('Feasible: ' + str(result.is_success()))\n",
    "\n",
    "  qt = result.GetSolution(q)\n",
    "  bt = result.GetSolution(b)\n",
    "  \n",
    "  line.set_xdata(qt[0,:])\n",
    "  line.set_ydata(qt[1,:])\n",
    "  start_pt.set_xdata(start_x)\n",
    "  start_pt.set_ydata(start_y)\n",
    "  goal_pt.set_xdata(goal_x)\n",
    "  goal_pt.set_ydata(goal_y)\n",
    "  for i in range(4):\n",
    "    blines[i].set_ydata(bt[i,:])\n",
    "  fig.canvas.draw()\n",
    "  \n",
    "interact(resolve, start_x=(-4,4,0.1), start_y=(-4,-1,0.1), goal_x=(-4,4,0.1), goal_y=(0,4,0.1));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
